use std::{iter::repeat_n, ops::Add};

use base64::{prelude::BASE64_STANDARD, Engine};

#[derive(Debug, Clone, Copy)]
pub struct ZucState {
    s: [U31; 16],
    r1: u32,
    r2: u32,
    x: [u32; 4],
}

/// An abstraction over GF(2 ** 31 - 1) used in ZUC algorithm
#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct U31(u32);

impl U31 {
    const MAX_U32: u32 = i32::MAX as u32;
    pub const MAX: U31 = U31(Self::MAX_U32);

    /// Create a new 31-bit word from a 32-bit word.
    ///
    /// U31::new(0) will become [`U31::MAX`], i.e. U31(2<sup>31</sup> - 1)
    #[inline]
    const fn new(x: u32) -> U31 {
        match x {
            0 => Self::MAX,
            _ => U31(x & Self::MAX_U32),
        }
    }

    /// High 16 bits of the 31-bit word
    #[inline]
    const fn hi(self) -> u16 {
        (self.0 >> 15) as u16
    }

    /// Low 16 bits of the 31-bit word
    #[inline]
    const fn lo(self) -> u16 {
        self.0 as u16
    }

    const fn rotate_left(self, bits: u32) -> U31 {
        let a = self.0;
        U31::new(a << bits | a >> (31 - bits))
    }
}

impl From<u32> for U31 {
    fn from(x: u32) -> Self {
        U31::new(x)
    }
}

impl Add for U31 {
    type Output = U31;

    /// Add mod (2<sup>31</sup> - 1).
    ///
    /// ZUC's GF(2<sup>31</sup> - 1) considers 0 and 2<sup>31</sup> - 1 to be equal, so it is normal
    /// for `U31(1) + U31((1 << 31) - 2)` to be `U31((1 << 31) - 1)`.
    fn add(self, rhs: Self) -> Self::Output {
        let res = self.0.wrapping_add(rhs.0);
        U31::new((res & Self::MAX_U32).wrapping_add(res >> 31))
    }
}

impl ZucState {
    fn init() -> ZucState {
        const D: [u16; 16] = [
            0b100010011010111,
            0b010011010111100,
            0b110001001101011,
            0b001001101011110,
            0b101011110001001,
            0b011010111100010,
            0b111000100110101,
            0b000100110101111,
            0b100110101111000,
            0b010111100010011,
            0b110101111000100,
            0b001101011110001,
            0b101111000100110,
            0b011110001001101,
            0b111100010011010,
            0b100011110101100,
        ];

        // See problem description.
        const K: [u8; 16] = [0; 16];
        const IV: [u8; 16] = [1; 16];

        // According to standard, `s[i] = K[i] || D[i] || IV[i]`, where `||` is digit concatenation.
        // Since K, D, IV are all const, we can generate S at compile time with a macro.
        macro_rules! const_s_expr {
            // Expand and move one expression to param 1.
            (@expand ([$($e:expr),*] [$i0:expr $(, $i:expr)* $(,)?])) => {
                const_s_expr!(@expand ([$($e,)* const_s_expr!(@map $i0)] [$($i),* ]))
            };
            // Nothing to expand, expand to param 1 directly.
            (@expand ([$($e:expr),*] [])) => {
                [$($e),*]
            };
            // Map an $i to U31
            (@map $i:expr) => {
                U31((K[$i] as u32) << 23 | (D[$i] as u32) << 8 | (IV[$i] as u32))
            };
            () => {
                const_s_expr!(@expand ([] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]))
            };
        }

        const S: [U31; 16] = const_s_expr!();

        let mut zuc = ZucState {
            s: S,
            r1: 0,
            r2: 0,
            x: [0; 4],
        };

        for _ in 0..32 {
            zuc.bit_reconstruction();
            let w = zuc.f();
            zuc.lfsr_with_initialization_mode(U31(w >> 1));
        }

        zuc
    }

    fn bit_reconstruction(&mut self) {
        /// Concat 2 16-bit words into a 32-bit word
        #[inline(always)]
        fn concat(x: u16, y: u16) -> u32 {
            (x as u32) << 16 | y as u32
        }

        let s = &self.s;
        self.x = [
            concat(s[15].hi(), s[14].lo()),
            concat(s[11].lo(), s[9].hi()),
            concat(s[7].lo(), s[5].hi()),
            concat(s[2].lo(), s[0].hi()),
        ]
    }

    fn f(&mut self) -> u32 {
        fn s_box(x: u32) -> u32 {
            const S0: [u8; 256] = [
                0x3E, 0x72, 0x5B, 0x47, 0xCA, 0xE0, 0x00, 0x33, 0x04, 0xD1, 0x54, 0x98, 0x09, 0xB9,
                0x6D, 0xCB, 0x7B, 0x1B, 0xF9, 0x32, 0xAF, 0x9D, 0x6A, 0xA5, 0xB8, 0x2D, 0xFC, 0x1D,
                0x08, 0x53, 0x03, 0x90, 0x4D, 0x4E, 0x84, 0x99, 0xE4, 0xCE, 0xD9, 0x91, 0xDD, 0xB6,
                0x85, 0x48, 0x8B, 0x29, 0x6E, 0xAC, 0xCD, 0xC1, 0xF8, 0x1E, 0x73, 0x43, 0x69, 0xC6,
                0xB5, 0xBD, 0xFD, 0x39, 0x63, 0x20, 0xD4, 0x38, 0x76, 0x7D, 0xB2, 0xA7, 0xCF, 0xED,
                0x57, 0xC5, 0xF3, 0x2C, 0xBB, 0x14, 0x21, 0x06, 0x55, 0x9B, 0xE3, 0xEF, 0x5E, 0x31,
                0x4F, 0x7F, 0x5A, 0xA4, 0x0D, 0x82, 0x51, 0x49, 0x5F, 0xBA, 0x58, 0x1C, 0x4A, 0x16,
                0xD5, 0x17, 0xA8, 0x92, 0x24, 0x1F, 0x8C, 0xFF, 0xD8, 0xAE, 0x2E, 0x01, 0xD3, 0xAD,
                0x3B, 0x4B, 0xDA, 0x46, 0xEB, 0xC9, 0xDE, 0x9A, 0x8F, 0x87, 0xD7, 0x3A, 0x80, 0x6F,
                0x2F, 0xC8, 0xB1, 0xB4, 0x37, 0xF7, 0x0A, 0x22, 0x13, 0x28, 0x7C, 0xCC, 0x3C, 0x89,
                0xC7, 0xC3, 0x96, 0x56, 0x07, 0xBF, 0x7E, 0xF0, 0x0B, 0x2B, 0x97, 0x52, 0x35, 0x41,
                0x79, 0x61, 0xA6, 0x4C, 0x10, 0xFE, 0xBC, 0x26, 0x95, 0x88, 0x8A, 0xB0, 0xA3, 0xFB,
                0xC0, 0x18, 0x94, 0xF2, 0xE1, 0xE5, 0xE9, 0x5D, 0xD0, 0xDC, 0x11, 0x66, 0x64, 0x5C,
                0xEC, 0x59, 0x42, 0x75, 0x12, 0xF5, 0x74, 0x9C, 0xAA, 0x23, 0x0E, 0x86, 0xAB, 0xBE,
                0x2A, 0x02, 0xE7, 0x67, 0xE6, 0x44, 0xA2, 0x6C, 0xC2, 0x93, 0x9F, 0xF1, 0xF6, 0xFA,
                0x36, 0xD2, 0x50, 0x68, 0x9E, 0x62, 0x71, 0x15, 0x3D, 0xD6, 0x40, 0xC4, 0xE2, 0x0F,
                0x8E, 0x83, 0x77, 0x6B, 0x25, 0x05, 0x3F, 0x0C, 0x30, 0xEA, 0x70, 0xB7, 0xA1, 0xE8,
                0xA9, 0x65, 0x8D, 0x27, 0x1A, 0xDB, 0x81, 0xB3, 0xA0, 0xF4, 0x45, 0x7A, 0x19, 0xDF,
                0xEE, 0x78, 0x34, 0x60,
            ];

            const S1: [u8; 256] = [
                0x55, 0xC2, 0x63, 0x71, 0x3B, 0xC8, 0x47, 0x86, 0x9F, 0x3C, 0xDA, 0x5B, 0x29, 0xAA,
                0xFD, 0x77, 0x8C, 0xC5, 0x94, 0x0C, 0xA6, 0x1A, 0x13, 0x00, 0xE3, 0xA8, 0x16, 0x72,
                0x40, 0xF9, 0xF8, 0x42, 0x44, 0x26, 0x68, 0x96, 0x81, 0xD9, 0x45, 0x3E, 0x10, 0x76,
                0xC6, 0xA7, 0x8B, 0x39, 0x43, 0xE1, 0x3A, 0xB5, 0x56, 0x2A, 0xC0, 0x6D, 0xB3, 0x05,
                0x22, 0x66, 0xBF, 0xDC, 0x0B, 0xFA, 0x62, 0x48, 0xDD, 0x20, 0x11, 0x06, 0x36, 0xC9,
                0xC1, 0xCF, 0xF6, 0x27, 0x52, 0xBB, 0x69, 0xF5, 0xD4, 0x87, 0x7F, 0x84, 0x4C, 0xD2,
                0x9C, 0x57, 0xA4, 0xBC, 0x4F, 0x9A, 0xDF, 0xFE, 0xD6, 0x8D, 0x7A, 0xEB, 0x2B, 0x53,
                0xD8, 0x5C, 0xA1, 0x14, 0x17, 0xFB, 0x23, 0xD5, 0x7D, 0x30, 0x67, 0x73, 0x08, 0x09,
                0xEE, 0xB7, 0x70, 0x3F, 0x61, 0xB2, 0x19, 0x8E, 0x4E, 0xE5, 0x4B, 0x93, 0x8F, 0x5D,
                0xDB, 0xA9, 0xAD, 0xF1, 0xAE, 0x2E, 0xCB, 0x0D, 0xFC, 0xF4, 0x2D, 0x46, 0x6E, 0x1D,
                0x97, 0xE8, 0xD1, 0xE9, 0x4D, 0x37, 0xA5, 0x75, 0x5E, 0x83, 0x9E, 0xAB, 0x82, 0x9D,
                0xB9, 0x1C, 0xE0, 0xCD, 0x49, 0x89, 0x01, 0xB6, 0xBD, 0x58, 0x24, 0xA2, 0x5F, 0x38,
                0x78, 0x99, 0x15, 0x90, 0x50, 0xB8, 0x95, 0xE4, 0xD0, 0x91, 0xC7, 0xCE, 0xED, 0x0F,
                0xB4, 0x6F, 0xA0, 0xCC, 0xF0, 0x02, 0x4A, 0x79, 0xC3, 0xDE, 0xA3, 0xEF, 0xEA, 0x51,
                0xE6, 0x6B, 0x18, 0xEC, 0x1B, 0x2C, 0x80, 0xF7, 0x74, 0xE7, 0xFF, 0x21, 0x5A, 0x6A,
                0x54, 0x1E, 0x41, 0x31, 0x92, 0x35, 0xC4, 0x33, 0x07, 0x0A, 0xBA, 0x7E, 0x0E, 0x34,
                0x88, 0xB1, 0x98, 0x7C, 0xF3, 0x3D, 0x60, 0x6C, 0x7B, 0xCA, 0xD3, 0x1F, 0x32, 0x65,
                0x04, 0x28, 0x64, 0xBE, 0x85, 0x9B, 0x2F, 0x59, 0x8A, 0xD7, 0xB0, 0x25, 0xAC, 0xAF,
                0x12, 0x03, 0xE2, 0xF2,
            ];

            u32::from_be_bytes([
                S0[(x >> 24) as usize],
                S1[(x >> 16 & 0xFF) as usize],
                S0[(x >> 8 & 0xFF) as usize],
                S1[(x & 0xFF) as usize],
            ])
        }

        fn l1(x: u32) -> u32 {
            x ^ x.rotate_left(2) ^ x.rotate_left(10) ^ x.rotate_left(18) ^ x.rotate_left(24)
        }

        fn l2(x: u32) -> u32 {
            x ^ x.rotate_left(8) ^ x.rotate_left(14) ^ x.rotate_left(22) ^ x.rotate_left(30)
        }

        let x = &mut self.x;
        let r1 = self.r1;
        let r2 = self.r2;

        let w = (x[0] ^ r1).wrapping_add(r2);

        let w1 = r1.wrapping_add(x[1]);
        let w2 = r2 ^ x[2];

        self.r1 = s_box(l1((w1 << 16) | (w2 >> 16)));
        self.r2 = s_box(l2((w2 << 16) | (w1 >> 16)));

        w
    }

    fn lfsr_with_initialization_mode(&mut self, u: U31) {
        let s = self.s;
        let v = s[15].rotate_left(15)
            + s[13].rotate_left(17)
            + s[10].rotate_left(21)
            + s[4].rotate_left(20)
            + s[0].rotate_left(8)
            + s[0];

        for i in 0..15 {
            self.s[i] = self.s[i + 1];
        }

        self.s[15] = u + v
    }

    fn lfsr_with_work_mode(&mut self) {
        let s = &self.s;
        let s16 = s[15].rotate_left(15)
            + s[13].rotate_left(17)
            + s[10].rotate_left(21)
            + s[4].rotate_left(20)
            + s[0].rotate_left(8)
            + s[0];

        for i in 0..15 {
            self.s[i] = self.s[i + 1];
        }

        self.s[15] = s16
    }

    fn generate_keystream(&mut self, n: usize) -> Vec<u32> {
        let mut keystream = Vec::with_capacity(n);

        self.bit_reconstruction();
        self.f();
        self.lfsr_with_work_mode();

        for _ in 0..n {
            self.bit_reconstruction();
            let z = self.f() ^ self.x[3];
            self.lfsr_with_work_mode();
            keystream.push(z);
        }
        keystream
    }

    fn encrypt(&mut self, input: &[u8]) -> Vec<u8> {
        let l = input.len();
        let key_stream = self.generate_keystream(l);

        key_stream
            .into_iter()
            .flat_map(|k| k.to_be_bytes())
            .zip(input)
            .map(|(ibs, k)| ibs ^ k)
            .collect()
    }
}

fn extend_pkcs7(input: &[u8]) -> Vec<u8> {
    let mut input = input.to_vec();
    let additional = 4 - input.len() % 4;
    input.extend(repeat_n(additional as u8, additional));
    input
}

pub fn encryption(input: String) -> String {
    let mut zuc = ZucState::init();

    let input = extend_pkcs7(input.as_bytes());

    BASE64_STANDARD.encode(zuc.encrypt(&input))
}
